'''
https://leetcode.cn/problems/count-univalue-subtrees/?envType=study-plan-v2&envId=premium-algo-100
'''
from typing import Optional


class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

class Solution:
    def countUnivalSubtrees(self, root: Optional[TreeNode]) -> int:
        if root is None:
            return 0
        count = 0
        def isUnivalSubtree(node: Optional[TreeNode]) -> bool:
            nonlocal count
            if node.left is None and node.right is None:
                count += 1
                return True
            isUnival = True
            if node.left is not None:
                isUnival = isUnivalSubtree(node.left) and isUnival and node.left.val == node.val
            if node.right is not None:
                isUnival = isUnivalSubtree(node.right) and isUnival and node.right.val == node.val
            if isUnival:
                count += 1
                return True
            else:
                return False
        isUnivalSubtree(root)
        return count
